
import numpy as np 

from vnpy.app.cta_strategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,
)

from vnpy.trader.constant import Direction, Offset

class MicroMartingaleStrategy(CtaTemplate):
    """
	马丁类策略的主要参数：
			加仓倍数（mul）
			加仓间距（dist）

	数字货币的高波动性：我们可以选择加仓间距例如3%

	平仓机制：盈利到一定百分比，全平。
    """

    author = "CQ"

    size = 1
    add_mul = 2.0
    add_dist = 0.03

    tp_pct = 0.02

    avg_price = 0

    long_signal = True
    long_entry = None
    long_entry_dt = None
    last_long_entry = None


    parameters = ["add_mul", "add_dist", "tp_pct"]
    variables = []

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        """"""
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)

        self.bg = BarGenerator(self.on_bar)
        self.am = ArrayManager()

        self.trade_list = []

        self.last_trade_size = 0

        self.first_open_price = 0
        self.count = 0

    def on_init(self):
        """
        Callback when strategy is inited.
        """
        self.write_log("策略初始化")
        self.load_bar(0)

    def on_start(self):
        """
        Callback when strategy is started.
        """
        self.write_log("策略启动")
        self.put_event()

    def on_stop(self):
        """
        Callback when strategy is stopped.
        """
        self.write_log("策略停止")

        self.put_event()

    def on_tick(self, tick: TickData):
        """
        Callback of new tick data update.
        """
        self.bg.update_tick(tick)

    def on_bar(self, bar: BarData):
        """
        Callback of new bar data update.
        """
        self.cancel_all()

        am = self.am
        am.update_bar(bar)
        if not am.inited:
            return

        if self.pos == 0:
            self.buy(bar.close_price, self.size)

            self.first_open_price = bar.close_price

        elif self.pos > 0:

            
            next_add_pos_price = self.first_open_price * (1 - self.add_dist * self.count)
            next_add_size = self.last_trade_size * self.add_mul

            avg_price = self.avg_calc()

            if (bar.close_price - avg_price) / avg_price > self.tp_pct:
                self.sell(bar.close_price, abs(self.pos))

            elif bar.close_price <= next_add_pos_price:
                

                if np.mean(self.am.volume[-30:]) < 600:
                    self.buy(next_add_pos_price, next_add_size)


        self.put_event()

    def on_order(self, order: OrderData):
        """
        Callback of new order data update.
        """
        pass

    def on_trade(self, trade: TradeData):
        """
        Callback of new trade data update.
        """
        if trade.offset == Offset.OPEN:
            if trade.direction == Direction.LONG:
                self.last_trade_size = trade.volume
                self.long_entry = trade.price
                self.long_entry_dt = trade.datetime

                self.trade_list.append((trade.price, trade.volume))

                self.count += 1
        
        elif trade.offset == Offset.CLOSE:
            if trade.direction == Direction.SHORT:

                self.trade_list = []
                self.last_trade_size = 0

                self.first_open_price = 0

                self.count = 0

        self.put_event()

    def on_stop_order(self, stop_order: StopOrder):
        """
        Callback of stop order update.
        """
        pass

    def avg_calc(self):
        """"""
        # 计算当前均价
        sum_price = 0
        sum_volume = 0

        for item in self.trade_list:
            _total_price = item[0] * item[1]
            sum_price += _total_price
            sum_volume += item[1]
        

        curr_avg_price = 0
        if sum_volume: 
            curr_avg_price = sum_price / sum_volume

        return curr_avg_price

    